package com.LeeCode;

/**
 * 二叉树染色
 */

public class CodeLCP034 {
    public static void main(String[] args) {

    }

    public int maxValue(TreeNode root, int k) {
        int[] dp = getDP(root, k);
        int ans = Integer.MIN_VALUE;
        for (int i = 0; i <= k; i++) {
            ans = Math.max(ans, dp[i]);
        }
        return ans;
    }

    public int[] getDP(TreeNode root, int maxCount) {
        // 当前节点为空时直接返回全为0的数组
        int[] dp = new int[maxCount + 1];
        if (root == null) return dp;

        // 获取左右子节点的dp结果
        int[] ldp = getDP(root.left, maxCount);
        int[] rdp = getDP(root.right, maxCount);

        // 当前节点不染色，最大值为左右子树染色最大值的和
        int lMax = Integer.MIN_VALUE;
        int rMax = Integer.MIN_VALUE;
        for (int i = 0; i <= maxCount; i++) {
            lMax = Math.max(lMax, ldp[i]);
            rMax = Math.max(rMax, rdp[i]);
        }
        dp[0] = lMax + rMax;

        // 当前节点染色个数为i时，取左右子节点染色个数和为(i-1)的所有情况的最大值
        for (int i = 1; i <= maxCount; i++) {
            for (int j = 0; j < i; j++) {
                dp[i] = Math.max(dp[i], root.val + ldp[j] + rdp[i - 1 - j]);
            }
        }
        return dp;
    }
}
